package stu

import (
	"context"
	"fmt"
	"github.com/gorilla/websocket"
	"log"
	"net/url"
	"sync"
	"time"
)

const (
	TextMessage   = websocket.TextMessage
	BinaryMessage = websocket.BinaryMessage
	CloseMessage  = websocket.CloseMessage
	PingMessage   = websocket.PingMessage
	PongMessage   = websocket.PongMessage
)

type WSConn struct {
	Conn      *websocket.Conn
	Mutex     sync.Mutex // 互斥锁，用于同步写操作
	ConnName  string
	ErrorChan chan error
	Addr      string
}

// --------------------------
// WebSocket 封装
// --------------------------
func NewWSConnection(addr string) (*WSConn, error) {
	var conn *websocket.Conn
	var _ error
	for {
		u, err := url.Parse(addr)
		if err != nil {
			return nil, fmt.Errorf("parse URL: %w", err)
		}
		log.Printf("[WS-CONN] %s", u.String())
		conn, _, err = websocket.DefaultDialer.Dial(u.String(), nil)
		if err == nil {
			return &WSConn{
				Conn:      conn,
				ConnName:  addr,
				Addr:      addr,
				ErrorChan: make(chan error),
				Mutex:     sync.Mutex{},
			}, nil
		}
		time.Sleep(10 * time.Second)
	}
}
func (w *WSConn) SendProgress(msg string) (string, error) {
	messageMutex := sync.Mutex{} // 用于同步 message 更新
	messageMutex.Lock()
	defer messageMutex.Unlock()
	err := w.SendMessage(TextMessage, msg)
	if err != nil {
		return "", fmt.Errorf("send progress error: %w", err)
	}
	recv, err := w.ReceiveMessages()
	if err != nil {
		return "", fmt.Errorf("receive progress error: %w", err)
	}
	return recv, nil
}
func (w *WSConn) SendMessage(messageType int, msg string) error {
	w.Mutex.Lock()
	defer w.Mutex.Unlock()
	// 设置合理的写超时
	if err := w.Conn.SetWriteDeadline(time.Now().Add(10 * time.Second)); err != nil {
		return fmt.Errorf("set write deadline failed: %w", err)
	}
	if err := w.Conn.WriteMessage(messageType, []byte(msg)); err != nil {
		return fmt.Errorf("write error: %w", err)
	}
	log.Printf("[WS-SEND-SUCCESS] %s", msg)
	return nil
}

func (w *WSConn) ReceiveMessages() (string, error) {
	err := w.Conn.SetReadDeadline(time.Now().Add(10 * time.Second))
	if err != nil {
		return "", err
	} // 延长读超时时间
	_, recv, err := w.Conn.ReadMessage()
	if err != nil {
		log.Printf("[WS-RECV-ERROR] ReadMessage error: %v", err)
		return "", err
	}
	log.Printf("[WS-RECV-ADDR] %s", w.Addr)
	log.Printf("[WS-RECV-CONNNAME] %s", w.ConnName)
	log.Printf("[WS-RECV-SUCCESS] %s", string(recv))
	return string(recv), nil
}

func (w *WSConn) Close() error {
	w.Mutex.Lock()
	defer w.Mutex.Unlock()
	err := w.Conn.Close()
	if err != nil {
		log.Printf("[WS-CLOSE] %v", err)
	}
	return err

}

func (w *WSConn) SendHeartbeat(wg *sync.WaitGroup, pingAddr string, ctx context.Context) {
	defer wg.Done()
	ticker := time.NewTicker(46 * time.Second) // 每46秒发送一次心跳
	defer ticker.Stop()
	for {
		select {
		case <-ticker.C:
			msg := "ping"
			err := w.SendMessage(PingMessage, msg)
			if err != nil {
				return
			}
		case <-ctx.Done():
			// 如果任务完成，退出心跳包协程
			log.Printf("[WS-PING] %s", pingAddr)
			return
		}
	}
}
